#!/usr/bin/env python
# coding=utf-8

# Copyright(C) 2022. Huawei Technologies Co.,Ltd. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import sys
import os
import enum
import numpy as np
import cv2
import mxpiOpenposeProto_pb2 as mxpiOpenposeProto
from StreamManagerApi import StreamManagerApi, MxDataInput, StringVector
sys.path.append("../proto")


COCO_PAIRS = [(1, 2), (1, 5), (2, 3), (3, 4), (5, 6), (6, 7), (1, 8), (8, 9), (9, 10), (1, 11),
             (11, 12), (12, 13), (1, 0), (0, 14), (14, 16), (0, 15), (15, 17), (2, 16), (5, 17)]  # = 19

COCO_PAIRS_RENDER = COCO_PAIRS[:-2]

COCO_COLORS = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],
              [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],
              [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]

if __name__ == '__main__':
    # init stream manager
    stream_manager_api = StreamManagerApi()
    ret = stream_manager_api.InitManager()
    if ret != 0:
        print("Failed to init Stream manager, ret=%s" % str(ret))
        exit()

    # create streams by pipeline config file
    with open("pipeline/Openpose.pipeline", "rb") as f:
        pipeline_str = f.read()
    ret = stream_manager_api.CreateMultipleStreams(pipeline_str)
    if ret != 0:
        print("Failed to create Stream, ret=%s" % str(ret))
        exit()
    data_input = MxDataInput()
    for i in range(1, 21):
        FILE_NAME = "./pic/test" + str(i) + ".jpg"
        try:
            with open(FILE_NAME, 'rb') as f:
                data_input.data = f.read()
        except FileNotFoundError:
            print("Test image", FILE_NAME, "doesn't exist. Exit.")
            exit()

        unique_id = stream_manager_api.SendData(b"classification+detection", 0, data_input)
        if unique_id < 0:
            print("Failed to send data to stream.")
            exit()
        keys = [b"mxpi_openposepostprocess0"]
        key_vec = StringVector()
        for key in keys:
            key_vec.push_back(key)          
        infer_result = stream_manager_api.GetProtobuf(b"classification+detection", 0, key_vec)
    #end code edit
        if infer_result.size() == 0:
            print("infer_result is null")
            exit()
        if infer_result[0].errorCode != 0:
            print("infer_result error. errorCode=%d" % (infer_result[0].errorCode))
            exit()
        result_personlist = mxpiOpenposeProto.MxpiPersonList()
        result_personlist.ParseFromString(infer_result[0].messageBuf)
        detect_person_list = result_personlist.personInfoVec
    # destroy streams
    stream_manager_api.DestroyAllStreams()
